/*

 * Copyright (c) 2017-2018, FastDev 刘强 (fastdev@163.com).

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package of.hotswapinject.injection;

import afu.org.checkerframework.checker.oigj.qual.O;
import of.hotswapinject.core.common.RuntimeContext;
import of.hotswapinject.core.common.Util;
import of.hotswapinject.core.plugin.*;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.AdviceAdapter;
import org.objectweb.asm.commons.GeneratorAdapter;
import org.objectweb.asm.commons.Method;

import java.util.ArrayList;
import java.util.List;

public class MethodHookAdvice extends AdviceAdapter {
    private String key;
    public final int access;
    public final Type returnType;
    public final Type[] argumentTypes;
    private TryCatchBuilder tryCatchBuilder = new TryCatchBuilder();

    public static final Type OBJECT_TYPE;

    public static final Type MethodMetaType;

    private static Method methodMetaConstructor;

    private static Method equalMethod;
    private static Method onEnterMethod;
    private static Method onExitMethod;

    private static Type enterPluginReturnParam;
    private static Type enterPluginReturnType;
    private static Type hookManagerType;
    private static Type objArrayType;
    private static Type exitPluginReturnParam;
    private static Type exitPluginReturnType;

    private class TryCatchBuilder {
        private List<Label> beginBlockList = new ArrayList<Label>();
        private List<Label> endBlockList = new ArrayList<Label>();
        private void addBegin() {
            Label begin = new Label();
            MethodHookAdvice.this.visitLabel(begin);
            beginBlockList.add(begin);
        }

        private void addEnd() {
            MethodHookAdvice.this.nop();
            Label end = new Label();
            MethodHookAdvice.this.visitLabel(end);
            endBlockList.add(end);
        }


        private void build() {
            /*
            Label handler = beginBlockList.get(beginBlockList.size() - 1);
            */
            MethodHookAdvice.this.nop();
            Label handler = new Label();
            MethodHookAdvice.this.visitLabel(handler);
            int tryCatchCount = 0;
            for(int lI = 0; lI < endBlockList.size(); lI++) {
                Label begin = beginBlockList.get(lI);
                Label end = endBlockList.get(lI);
                visitTryCatchBlock(begin, end, handler, null);
                tryCatchCount++;
            }
            if(tryCatchCount > 0) {
                injectMethodExit(Opcodes.ATHROW);
                mv.visitInsn(Opcodes.ATHROW);
            }
        }
    }

    private static class ObjectAsmBuilder {
        private MethodHookAdvice advice = null;
        private Type type = null;
        private ObjectAsmBuilder() {

        }
        public static ObjectAsmBuilder newInstance(MethodHookAdvice advice, Type type, Method constructor) {
            ObjectAsmBuilder builder = new ObjectAsmBuilder();
            builder.advice = advice;
            builder.type = type;
            advice.newInstance(type);
            advice.dup();
            advice.invokeConstructor(type, constructor);
            return builder;
        }

        public ObjectAsmBuilder setFiled(String fieldName, boolean fieldValue) {
            advice.dup();
            advice.push(fieldValue);
            advice.putField(this.type, fieldName, Type.BOOLEAN_TYPE);
            return this;
        }

        public ObjectAsmBuilder setFiled(String fieldName, int fieldValue) {
            advice.dup();
            advice.push(fieldValue);
            advice.putField(this.type, fieldName, Type.INT_TYPE);
            return this;
        }
    }

    static {
        OBJECT_TYPE = (Type) Util.getFieldValue(null, GeneratorAdapter.class, "OBJECT_TYPE");

        MethodMetaType = Type.getType(MethodMeta.class);
        try {
            methodMetaConstructor = Method.getMethod(MethodMeta.class.getConstructor());
        } catch (NoSuchMethodException e) {
            e.printStackTrace();
            throw new RuntimeException("MethodMeta have no default Constructor");
        }

        try {
            java.lang.reflect.Method equalMethodInfo = Object.class.getDeclaredMethod("equals", Object.class);
            equalMethod = Method.getMethod(equalMethodInfo);
        } catch (NoSuchMethodException e) {
            e.printStackTrace();
            throw new RuntimeException("equalMethod == null", e);
        }
        try {
            java.lang.reflect.Method onEnterMethodInfo = MethodHook.class.getDeclaredMethod("onEnter", MethodMeta.class, String.class, Object[].class);
            onEnterMethod = Method.getMethod(onEnterMethodInfo);
        } catch (NoSuchMethodException e) {
            e.printStackTrace();
            throw new RuntimeException("MethodHook.onEnter error", e);
        }

        try {
            java.lang.reflect.Method onExitMethodInfo = MethodHook.class.getDeclaredMethod("onExit", Object.class, MethodMeta.class, String.class, int.class, Object[].class);
            onExitMethod = Method.getMethod(onExitMethodInfo);
        } catch (NoSuchMethodException e) {
            e.printStackTrace();
            throw new RuntimeException("MethodHook.onExit error", e);
        }

        enterPluginReturnParam = Type.getType(IMethodEnterPlugin.ReturnParam.class);
        enterPluginReturnType = Type.getType(IMethodEnterPlugin.ReturnTypeEnum.class);
        hookManagerType = Type.getType(MethodHook.class);
        objArrayType = Type.getType(Object[].class);
        exitPluginReturnParam = Type.getType(IMethodExitPlugin.ReturnParam.class);
        exitPluginReturnType = Type.getType(IMethodExitPlugin.ReturnTypeEnum.class);
    }

    protected MethodHookAdvice(String key, int var1, MethodVisitor var2, int var3, String var4, String var5) {
        super(var1, var2, var3, var4, var5);
        this.key = key;
        this.access = (Integer) Util.getFieldValue(this, GeneratorAdapter.class, "access");
        this.argumentTypes = (Type[])Util.getFieldValue(this, GeneratorAdapter.class, "argumentTypes");
        this.returnType = (Type)Util.getFieldValue(this, GeneratorAdapter.class, "returnType");
    }

    public void nop() {
        this.mv.visitInsn(Opcodes.NOP);
    }

    public void dup(Type type) {
        if(Type.LONG_TYPE.equals(type)
                || Type.DOUBLE_TYPE.equals(type)) {
            this.mv.visitInsn(Opcodes.DUP2);

        }
        else {
            this.mv.visitInsn(Opcodes.DUP);
        }
    }

    public void loadInParams() {
        if ((this.access & 8) == 0) {
            this.push(this.argumentTypes.length + 1);
            this.newArray(OBJECT_TYPE);

            this.dup();
            this.push(0);
            this.mv.visitVarInsn(Opcodes.ALOAD, 0);
            this.arrayStore(OBJECT_TYPE);

            for(int var1 = 0; var1 < this.argumentTypes.length; ++var1) {
                this.dup();
                this.push(var1 + 1);
                this.loadArg(var1);
                this.box(this.argumentTypes[var1]);
                this.arrayStore(OBJECT_TYPE);
            }
        }
        else {
            this.loadArgArray();
        }
    }

    public boolean hasArgs() {
        return this.argumentTypes.length > 0;
    }

    public boolean hasReturnValue() {
        return !Type.VOID_TYPE.equals(this.returnType);
    }

    public boolean isStatic() {
        if ((this.access & 8) == 0) {
            return false;
        }
        return true;
    }

    private void pushMethodMeta() {
        ObjectAsmBuilder.newInstance(this, MethodMetaType, methodMetaConstructor)
                .setFiled("hasArgs", hasArgs())
                .setFiled("hasReturnValue", hasReturnValue())
                .setFiled("isStatic", isStatic());
    }

    @Override
    public void onMethodEnter() {
        super.onMethodEnter();
        pushMethodMeta();
        push(key);
        this.loadInParams();
        invokeStatic(hookManagerType, onEnterMethod);
        dup();
        getField(enterPluginReturnParam, "ReturnType", enterPluginReturnType);
        getStatic(enterPluginReturnType, "ReturnResult", enterPluginReturnType);
        invokeVirtual(enterPluginReturnType, equalMethod);
        Label notReturnType = new Label();
        ifZCmp(Opcodes.IFEQ, notReturnType);
        if(Type.VOID_TYPE.equals(this.returnType)) {
            pop();
            returnValue();
        }
        else {
            getField(enterPluginReturnParam, "Value", OBJECT_TYPE);
            unbox(this.returnType);
            returnValue();
        }

        visitLabel(notReturnType);
        dup();
        getField(enterPluginReturnParam, "ReturnType", enterPluginReturnType);
        getStatic(enterPluginReturnType, "ReplaceArgs", enterPluginReturnType);
        invokeVirtual(enterPluginReturnType, equalMethod);
        Label notReplaceType = new Label();
        ifZCmp(Opcodes.IFEQ, notReplaceType);
        getField(enterPluginReturnParam, "Value", OBJECT_TYPE);
        unbox(objArrayType);
        for(int i1 = 0; i1 < this.argumentTypes.length; i1++) {
            Type argType = this.argumentTypes[i1];
            dup();
            push(i1);
            arrayLoad(OBJECT_TYPE);
            unbox(argType);
            storeArg(i1);
        }

        visitLabel(notReplaceType);
        pop();
        tryCatchBuilder.addBegin();
    }

    @Override
    public void onMethodExit(int opcode) {
        super.onMethodExit(opcode);
        if(opcode == Opcodes.ATHROW) {
            return;
        }

        injectMethodExit(opcode);
    }

    @Override
    public void visitInsn(int opcode) {
        boolean isReturn = (opcode == Opcodes.IRETURN
                || opcode == Opcodes.LRETURN
                || opcode == Opcodes.FRETURN
                || opcode == Opcodes.DRETURN
                || opcode == Opcodes.ARETURN
                || opcode == Opcodes.RETURN);

        if(isReturn) {
            tryCatchBuilder.addEnd();
        }

        super.visitInsn(opcode);

        if(isReturn) {
            tryCatchBuilder.addBegin();
        }
    }

    private void injectMethodExit(int opcode) {
        if(Type.VOID_TYPE.equals(this.returnType)) {
            this.mv.visitInsn(Opcodes.ACONST_NULL);
        }
        else if(opcode == Opcodes.ATHROW){
            this.mv.visitInsn(Opcodes.ACONST_NULL);
        }
        else {
            dup(this.returnType);
            box(this.returnType);
        }
        pushMethodMeta();
        push(key);
        push(opcode);
        this.loadInParams();
        invokeStatic(hookManagerType, onExitMethod);
        dup();
        getField(exitPluginReturnParam, "ReturnType", exitPluginReturnType);
        getStatic(exitPluginReturnType, "ReturnResult", exitPluginReturnType);
        invokeVirtual(exitPluginReturnType, equalMethod);
        Label notReturnType = new Label();
        ifZCmp(Opcodes.IFEQ, notReturnType);
        if(Type.VOID_TYPE.equals(this.returnType)) {
            pop();
            returnValue();
        }
        else {
            getField(exitPluginReturnParam, "Value", OBJECT_TYPE);
            unbox(this.returnType);
            returnValue();
        }
        visitLabel(notReturnType);
        pop();
    }


    @Override
    public void visitMaxs(int var1, int var2) {
        tryCatchBuilder.build();
        super.visitMaxs(var1, var2);
    }
}
